# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
import random
import torch
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


def random_factor_pair(M):
    # 找出所有正因子
    factors = [i for i in range(1, M+1) if M % i == 0]
    # 随机选择一个因子A
    A = random.choice(factors)
    B = M // A
    return A, B

@GENERATOR_REGISTRY.register("ascend_generate_grouped_matmul_swiglu_quant")
class AscendGroupedMatmulSwigluQuant(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        return input_case

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        x = case_config.inputs[0]
        weight = case_config.inputs[1]
        weightScale = case_config.inputs[4]
        xScale = case_config.inputs[5]
        groupList = case_config.inputs[6]

        M = x.shape[0]

        #K需要是16，N需要是32
        weight.shape[3],weight.shape[4] = 16,32
        E,K,N= weight.shape[0], weight.shape[2] * weight.shape[3],weight.shape[1] * weight.shape[4]
        
        #K轴不能大于等于65536
        if K >= 65536:
            weight.shape[2] = 4095
            weight.shape[3] = 16
        x.shape[1] = weight.shape[2] * weight.shape[3]
        #N轴不能大于8192
        if N > 8192:
            fac_1,fac_4 = random_factor_pair(4096*2)
            weight.shape[1],weight.shape[4] = fac_1,fac_4
            N = 8192
        weightScale.shape[0] = E
        weightScale.shape[1] = N
        xScale.shape[0] = M
        # print(case_config)
        return case_config